package com.element.oauth2.server.global.provider;

import com.element.oauth2.server.global.generator.OAuthAuthorizationCodeGenerator;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.security.authentication.AnonymousAuthenticationToken;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.oauth2.server.authorization.*;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationCodeRequestAuthenticationException;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationCodeRequestAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationConsentAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

public final class OAuthConsentRequestAuthenticationProvider implements AuthenticationProvider {

    private static final OAuth2TokenType STATE_TOKEN_TYPE = new OAuth2TokenType("state");
    private final Log logger = LogFactory.getLog(this.getClass());
    private final RegisteredClientRepository registeredClientRepository;
    private final OAuth2AuthorizationService authorizationService;
    private final OAuth2AuthorizationConsentService authorizationConsentService;
    private final OAuth2TokenGenerator<OAuth2AuthorizationCode> authorizationCodeGenerator = new OAuthAuthorizationCodeGenerator();

    public OAuthConsentRequestAuthenticationProvider(RegisteredClientRepository registeredClientRepository, OAuth2AuthorizationService authorizationService, OAuth2AuthorizationConsentService authorizationConsentService) {
        Assert.notNull(registeredClientRepository, "registeredClientRepository cannot be null");
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        Assert.notNull(authorizationConsentService, "authorizationConsentService cannot be null");
        this.registeredClientRepository = registeredClientRepository;
        this.authorizationService = authorizationService;
        this.authorizationConsentService = authorizationConsentService;
    }

    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        OAuth2AuthorizationConsentAuthenticationToken authorizationConsentAuthentication = (OAuth2AuthorizationConsentAuthenticationToken) authentication;
        OAuth2Authorization authorization = this.authorizationService.findByToken(authorizationConsentAuthentication.getState(), STATE_TOKEN_TYPE);
        if (authorization == null) {
            throwError("invalid_request", "authorization", authorizationConsentAuthentication, null, null);
            return null;
        }
        Authentication principal = (Authentication) authorizationConsentAuthentication.getPrincipal();
        if (!isPrincipalAuthenticated(principal) || !principal.getName().equals(authorization.getPrincipalName())) {
            throwError("invalid_request", "principal", authorizationConsentAuthentication, null, null);
            return null;
        }

        RegisteredClient registeredClient = this.registeredClientRepository.findByClientId(authorizationConsentAuthentication.getClientId());
        if (null == registeredClient) {
            throwError("invalid_request", "client", authorizationConsentAuthentication, registeredClient, null);
            return null;
        }
        if (!authorization.getRegisteredClientId().equals(registeredClient.getId())) {
            throwError("invalid_request", "client", authorizationConsentAuthentication, registeredClient, null);
        }
        if (!CollectionUtils.isEmpty(registeredClient.getScopes()) && !registeredClient.getScopes().containsAll(authorizationConsentAuthentication.getScopes())) {
            throwError("invalid_request", "scopes", authorizationConsentAuthentication, registeredClient, null);
        }
        OAuth2AuthorizationRequest authorizationRequest = authorization.getAttribute(OAuth2AuthorizationRequest.class.getName());
        Set<String> authorizedScopes = new HashSet<>(authorizationConsentAuthentication.getScopes());
        if (this.logger.isTraceEnabled()) {
            this.logger.trace("Validated authorization consent request parameters");
        }
        OAuth2AuthorizationConsent currentAuthorizationConsent = this.authorizationConsentService.findById(authorization.getRegisteredClientId(), authorization.getPrincipalName());
        Set<String> currentAuthorizedScopes = currentAuthorizationConsent != null ? currentAuthorizationConsent.getScopes() : Collections.emptySet();
        if (!currentAuthorizedScopes.isEmpty()) {
            for (String requestedScope : authorizationConsentAuthentication.getScopes()) {
                if (currentAuthorizedScopes.contains(requestedScope)) {
                    authorizedScopes.add(requestedScope);
                }
            }
        }
        if (!authorizedScopes.isEmpty() && authorizationConsentAuthentication.getScopes().contains("open_id")) {
            authorizedScopes.add("open_id");
        }

        OAuth2AuthorizationConsent.Builder authorizationConsentBuilder;
        if (currentAuthorizationConsent != null) {
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Retrieved existing authorization consent");
            }
            authorizationConsentBuilder = OAuth2AuthorizationConsent.from(currentAuthorizationConsent);
        } else {
            authorizationConsentBuilder = OAuth2AuthorizationConsent.withId(authorization.getRegisteredClientId(), authorization.getPrincipalName());
        }
        if (!CollectionUtils.isEmpty(registeredClient.getScopes()) &&
                CollectionUtils.isEmpty(authorizationConsentAuthentication.getScopes())
                && !registeredClient.getScopes().containsAll(authorizationConsentAuthentication.getScopes())) {
            throwError("access_denied", "scopes", authorizationConsentAuthentication, registeredClient, authorizationRequest);
            return null;
        }
        if (!CollectionUtils.isEmpty(registeredClient.getScopes())) {
            authorizedScopes.forEach(authorizationConsentBuilder::scope);
            Set<GrantedAuthority> authorities = new HashSet<>();
            authorizationConsentBuilder.authorities(authorities::addAll);
            if (authorities.isEmpty()) {
                throwError("access_denied", "scopes", authorizationConsentAuthentication, registeredClient, authorizationRequest);
                return null;
            }
        } else {
            authorizationConsentBuilder.authorities(authorities -> authorities.add(new SimpleGrantedAuthority("SCOPE_NONE")));
        }

//        if (currentAuthorizationConsent != null) {
//            this.authorizationConsentService.remove(currentAuthorizationConsent);
//            this.authorizationService.remove(authorization);
//        }

        OAuth2AuthorizationConsent authorizationConsent = authorizationConsentBuilder.build();
        if (!authorizationConsent.equals(currentAuthorizationConsent)) {
            this.authorizationConsentService.save(authorizationConsent);
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Saved authorization consent");
            }
        }

        OAuth2TokenContext tokenContext = createAuthorizationCodeTokenContext(authorizationConsentAuthentication, registeredClient, authorization, authorizedScopes);
        OAuth2AuthorizationCode authorizationCode = this.authorizationCodeGenerator.generate(tokenContext);
        if (authorizationCode == null) {
            OAuth2Error error = new OAuth2Error("server_error", "The token generator failed to generate the authorization code.", "https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1");
            throw new OAuth2AuthorizationCodeRequestAuthenticationException(error, (OAuth2AuthorizationCodeRequestAuthenticationToken) null);
        } else {
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Generated authorization code");
            }
            OAuth2Authorization updatedAuthorization = OAuth2Authorization.from(authorization).authorizedScopes(authorizedScopes).token(authorizationCode).attributes((attrs) -> {
                attrs.remove("state");
            }).build();
            this.authorizationService.save(updatedAuthorization);
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Saved authorization");
            }
            assert authorizationRequest != null;
            String redirectUri = authorizationRequest.getRedirectUri();
            if (!StringUtils.hasText(redirectUri)) {
                redirectUri = registeredClient.getRedirectUris().iterator().next();
            }

            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Authenticated authorization consent request");
            }

            return new OAuth2AuthorizationCodeRequestAuthenticationToken(authorizationRequest.getAuthorizationUri(), registeredClient.getClientId(), principal, authorizationCode, redirectUri, authorizationRequest.getState(), authorizedScopes);
        }
    }

    public boolean supports(Class<?> authentication) {
        return OAuth2AuthorizationConsentAuthenticationToken.class.isAssignableFrom(authentication);
    }

    private static OAuth2TokenContext createAuthorizationCodeTokenContext(OAuth2AuthorizationConsentAuthenticationToken authorizationConsentAuthentication, RegisteredClient registeredClient, OAuth2Authorization authorization, Set<String> authorizedScopes) {
        return DefaultOAuth2TokenContext.builder().registeredClient(registeredClient).principal((Authentication) authorizationConsentAuthentication.getPrincipal()).authorization(authorization).authorizationServerContext(AuthorizationServerContextHolder.getContext()).tokenType(new OAuth2TokenType("code")).authorizedScopes(authorizedScopes).authorizationGrantType(AuthorizationGrantType.AUTHORIZATION_CODE).authorizationGrant(authorizationConsentAuthentication).build();
    }

    private static boolean isPrincipalAuthenticated(Authentication principal) {
        return principal != null && !AnonymousAuthenticationToken.class.isAssignableFrom(principal.getClass()) && principal.isAuthenticated();
    }

    private static void throwError(String errorCode, String parameterName, OAuth2AuthorizationConsentAuthenticationToken authorizationConsentAuthentication, RegisteredClient registeredClient, OAuth2AuthorizationRequest authorizationRequest) {
        OAuth2Error error = new OAuth2Error(errorCode, "OAuth 2.0 Parameter: " + parameterName, "https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1");
        throwError(error, parameterName, authorizationConsentAuthentication, registeredClient, authorizationRequest);
    }

    private static void throwError(OAuth2Error error, String parameterName, OAuth2AuthorizationConsentAuthenticationToken authorizationConsentAuthentication, RegisteredClient registeredClient, OAuth2AuthorizationRequest authorizationRequest) {
        String redirectUri = resolveRedirectUri(authorizationRequest, registeredClient);
        if (error.getErrorCode().equals("invalid_request") && (parameterName.equals("client_id") || parameterName.equals("state"))) {
            redirectUri = null;
        }

        String state = authorizationRequest != null ? authorizationRequest.getState() : authorizationConsentAuthentication.getState();
        Set<String> requestedScopes = authorizationRequest != null ? authorizationRequest.getScopes() : authorizationConsentAuthentication.getScopes();
        OAuth2AuthorizationCodeRequestAuthenticationToken authorizationCodeRequestAuthenticationResult = new OAuth2AuthorizationCodeRequestAuthenticationToken(authorizationConsentAuthentication.getAuthorizationUri(), authorizationConsentAuthentication.getClientId(), (Authentication) authorizationConsentAuthentication.getPrincipal(), redirectUri, state, requestedScopes, null);
        throw new OAuth2AuthorizationCodeRequestAuthenticationException(error, authorizationCodeRequestAuthenticationResult);
    }

    private static String resolveRedirectUri(OAuth2AuthorizationRequest authorizationRequest, RegisteredClient registeredClient) {
        if (authorizationRequest != null && StringUtils.hasText(authorizationRequest.getRedirectUri())) {
            return authorizationRequest.getRedirectUri();
        } else {
            return registeredClient != null ? registeredClient.getRedirectUris().iterator().next() : null;
        }
    }
}